# -*- coding: utf-8 -*-

__author__="Boris Avdeev"
__date__ ="$Mar 8, 2011 10:09:48 AM$"


from TCdata import TCdata

#TODO: Write unit tests.

import numpy as np
import pymc as pm

c_aft = {
    'lam_d': 1.55125e-10,  # 1/y [Tagami, 2005]
    #'lam_d': 1.55125e-4,  # 1/My [Tagami, 2005]
#    'g'    : 0.5              
    }

def aft_age(Ns, A, P, zeta, g, fix=0.375):
    """ Compute AFT age for ICP-MS data. 
    Standard age equation underestimates age for small counts.
    Adding 3/8 to Ns is a dirty fix [Anscombe 1948].
    g is 0.5, but should be 1 if Donelick's new zeta is used."""
    return 1.0e-6 / c_aft['lam_d'] * np.log(1.+ c_aft['lam_d'] * zeta * g * (Ns + fix) / A / P)

def aft_exp_ns(t, A, P, zeta, g):
    """Expected Ns count (Inverse of the AFT age eq.)
    g is 0.5, but should be 1 if Donelick's new zeta is used."""
    return A*P*(np.exp(c_aft['lam_d'] * np.array(t)*1.0e6)-1)/c_aft['lam_d']/zeta/g

class AFTAgeModel(object):
    """Single grain and sample age model"""
    def __init__(self, aft, sample_age, P1s=0.1, Age=None):
        """If Age or P1s are None, they are modelled. Alternatively, constants or extermal models can be passed."""
#        tries = 1000
        max_age=1000. # Limit of AFT
        name=aft.name
        Ns = aft.counts['Ns']
        P  = aft.counts['P']
        A  = aft.counts['A']
#        while tries > 0:
#            try:
        # PRIORS
        if P1s is None: self.P1s = pm.Uniform("P1s_%s"%name, 0, 0.1, value=0.05)
        elif isinstance(P1s,float):
            self.P1s = P * P1s # 10% is what Ray finds from Durango replicates.
            #self.P1s[self.P1s<0.05] = 0.05  # 0.05 was crudely estimated from Caucasus bedrock samples.
        else: self.P1s = P1s
        self.P =    pm.Uniform("P_%s"    % name, 0., 10., value = P, plot=False, trace=False)
        self.Zeta = pm.Uniform("Zeta_%s" % name, 0., 50., value = aft.zeta[0], plot=False, trace=False)

        if Age is None:
            if sample_age: # assume common age for all grains
                init_age = aft_age(np.sum(Ns), 1., np.sum(A*P), aft.zeta[0], aft.g)
            else: # each grain has individual age
                init_age = aft_age(Ns, A, P, aft.zeta[0], aft.g)
            self.Age = pm.Uniform("AFTAge_%s"  % name, 0., max_age, value = init_age)
        else: self.Age = Age

        # DATA
        self.ObsP = pm.Normal("ObsP_%s" % name,  mu = self.P, tau = 1./self.P1s**2, value = P, observed=True)
        self.ObsZeta = pm.Normal("ObsZeta_%s" % name, mu = self.Zeta, tau = 1./aft.zeta[1]**2, value=aft.zeta[0], observed=True)
        self.ExpNs = pm.Lambda("ExpNs_%s"%name,
            lambda t=self.Age, a=A, p=self.P, zeta=self.Zeta, g=aft.g : aft_exp_ns(t,a,p,zeta,g),plot=False)
        self.ObsNs = pm.Poisson("Ns_%s"%name, mu = self.ExpNs, value = Ns, observed=True)

        # SIMULATION
        self.SimNs = pm.Deterministic(eval=pm.rpoisson,name="SimNs_%s"%name,doc='', parents={'mu':self.ExpNs},plot=False)
#               break
#           except pm.ZeroProbability:
#               tries-=1
#       if tries==0: raise pm.ZeroProbability

    def get_list(self):
        return self.__dict__.values()
        


class AFTAnnealModel(object):
    def __init__(self, aft, t, T, P1s=0.1):
        name=aft.name
        Ns = aft.counts['Ns']
        P  = aft.counts['P']
        A  = aft.counts['A']
        len = aft.lengths['length']
        ang = aft.lengths['angle']
        dpar= aft.lengths['Dpar']
        
        cproj=True
        # PRIORS
        if P1s is None: self.P1s = pm.Uniform("P1s_%s"%name, 0, 0.1, value=0.05)
        elif isinstance(P1s,float):
            self.P1s = P * P1s # 10% is what Ray finds from Durango replicates.
        else: self.P1s = P1s
        self.P =    pm.Uniform("P_%s"    % name, 0., 10., value = P, plot=False,trace=False)
        self.Zeta = pm.Uniform("Zeta_%s" % name, 0., 50., value = aft.zeta[0], plot=False,trace=False)
        
        # DATA
        self.ObsP = pm.Normal("ObsP_%s" % name,  mu = self.P, tau = 1./self.P1s**2, value = P, observed=True)
        self.ObsZeta = pm.Normal("ObsZeta_%s" % name, mu = self.Zeta, tau = 1./aft.zeta[1]**2, value=aft.zeta[0], observed=True)
        self.ExpNs = pm.Lambda("ExpNs_%s"%name,
            lambda t=self.Age, a=A, p=self.P, zeta=self.Zeta, g=aft.g : aft_exp_ns(t,a,p,zeta,g),plot=False)
        self.ObsNs = pm.Poisson("Ns_%s"%name, mu = self.ExpNs, value = Ns, observed=True)

        # SIMULATION
        self.SimNs = pm.Deterministic(eval=pm.rpoisson,name="SimNs_%s"%name,doc='', parents={'mu':self.ExpNs},plot=False)
            

    def lsig(self, cproj=True):
        """Standard deviation (?) of a track length""" 
        if cproj:
            l = l0 * r
            return 0.01*l*l - 0.2827*l + 2.501
        else:
            #h is computed according to Lutz1991, app
            if r > 0.68: h = 1.32-0.47*r  # 1.39-.61*r
            elif r < 0.43: h = 3          # 2.53
            else: h = 6.5-8.1*r           # 5.08-5.93*r
        

        
  
    def l0(self,Dpar=None,cproj=False):
        """Initial track length [HeFTy Preferences]"""
        if Dpar is None or 0: return 16.3
        if cproj: return 0.287 * Dpar + 15.582
        return 0.258 * Dpar + 15.391



    def get_list(self):
        return self.__dict__.values()
        





class AFTdata(TCdata):
    """Class for storing and analyzing apatite fission-track ICP-MS data.
    Length data is not used at the moment."""
        
    def __init__(self, cnt_fname, len_fname, new_zeta, sample=None):
        """ Reads AFTSolve data files """
        TCdata.__init__(self)
        self.method = 'AFT'
        cnt_file = open(cnt_fname)
        cnt_head = cnt_file.readlines(6)
        
        cnt_file.seek(0)
        if sample is None:
            sample = cnt_head[1].strip()
        self.name=sample
       
        zeta_line = cnt_head[4].split() 
        self.zeta   = (float(zeta_line[0]), float(zeta_line[1]))
        if new_zeta: self.g = 1
        else: self.g=0.5
        
        counts = np.genfromtxt(cnt_file,skip_header=6,dtype=None,names="Ns,A,P,P_1s,Dpar")
        self.counts = counts
        cnt_file.close()

        try:
            len_file = open(len_fname)
#            len_head = len_file.readlines(3)
            len_file.seek(0)
#            if(sample!=len_head[0].strip()):
#                message = "WARNING: Mismatch of sample names between %s and %s files" % (cnt_file.name, len_file.name)
#                #raise Exception, message
#                print message
            self.lengths = np.genfromtxt(len_file,skip_header=2, dtype=None, names=True)
            len_file.close()
            self.n_len = len(self.lengths['length'])
        except IOError:
            print "No length data for sample %s." % sample

        n_cnt = len(self.counts['Ns'])
        self.n_cnt=n_cnt
        
        #-------INITIATE MODELS-------------------------------------------------
        try: self.sa_model = AFTAgeModel(self, True)
        except pm.ZeroProbability: print "Cannot initiate SA model"
        try: self.sga_model = AFTAgeModel(self, False)
        except pm.ZeroProbability: print "Cannot initiate SGA model"
        
        
        
        

    @property
    def __ni(self):
        """Estimate induced track densities (required by QTQt)
        from LA-ICP-MS P value (P = U/Ca, Donelick et al., 2005)"""
        zeta  = self.zeta * 20  # Number to approximately match real value.  Should not matter without error handling.
        rho_d = 1000000         # Number to approximately match real counts. Should not matter without error handling.
        #n_d = 1000             # Number to approximately match real counts. Should not matter without error handling.
        return self.counts['A'] * self.counts['P'] * (rho_d * zeta / self.zeta)


    @property
    def quick_ages(self):
        """ Vector of single-grain ages."""
        return aft_age(self.counts['Ns'], self.counts['A'], self.counts['P'], self.zeta[0], self.g)

    @property
    def naive_ages(self):
        """ Vector of single-grain ages (without adding 3/8 to Ns)."""
        return aft_age(self.counts['Ns'], self.counts['A'], self.counts['P'], self.zeta[0],self.g,fix=0)

    @property
    def quick_ages_errs(self):
        """ Vector of single-grain age errors.
        WARNING: Do not use this. Poor estimate, esp. for small Ns."""
        qa = self.quick_ages
        ns = self.counts['Ns'].copy()
        #Donelick (e-mail)
        ns[ns==0] = 1.            
        qe = qa * (1/ns+
                (self.counts['P_1s']/self.counts['P'])**2+   # Donelick sets this to 0.10
                (self.zeta[1]/self.zeta[0])**2)**0.5
        return zip(qa,qe)

       
    @property
    def pooled_age(self):
        """Compute pooled age and error [e.g. Galbraith, 2005]"""
        Ns = self.counts['Ns']
        A = self.counts['A']
        P = self.counts['P']
        P1s=self.counts['P_1s']
        age = aft_age( np.sum(Ns), 1., np.sum(A*P), self.zeta[0], self.g)
        age1s_rel = (1./np.sum(Ns) + 
                    np.sum(P1s**2 * A**2) / np.sum(A*P)**2 + 
                    (self.zeta[1]/self.zeta[0])**2      )**0.5
        return age, age * age1s_rel
    
    
    #@property
    def sa_discrepancy(self, plot=False):
        M = self.sa_mcmc()
        D_obs, D_sim = pm.discrepancy(self.counts['Ns'],M.SimNs,M.ExpNs)
        if plot:
            pm.Matplot.discrepancy_plot((D_obs,D_sim), name="%s_sa" % self.name, report_p=True)
        
        count = sum(s>o for o,s in zip(D_obs,D_sim))
        return 1.*count/len(D_obs)
    
    #@property
    def sa_discplot(self):
        M = self.sa_mcmc()
        D_obs, D_sim = pm.discrepancy(self.counts['Ns'],M.SimNs,M.ExpNs)
        return pm.Matplot.discrepancy_plot((D_obs,D_sim), name="%s_sa" % self.name, report_p=True)

    def sga_discrepancy(self, plot=False):
        """Check whether sg_ages fit sample_age model"""
        M = self.sga_mcmc()
        D_obs, D_sim = pm.discrepancy(self.counts['Ns'],M.SimNs.trace(),M.ExpNs.trace())
        if plot: pm.Matplot.discrepancy_plot((D_obs,D_sim), name="%s_sga" % self.name, report_p=True)
        count = sum(s>o for o,s in zip(D_obs,D_sim))
        return 1.*count/len(D_obs)
                        
    
    def sga2csv(self,fname=None):
        """Output single grain age data"""
        if fname==None: fname="%s_sg.csv" % self.name
        f = open(fname, 'w')
        qages = self.quick_ages_errs
        sgages = self.sga
        f.write("Ns, A, P, P_1s, Dpar, Age, Age1s, Age (MCMC), Age1s (MCMC), HPD_lo, HPD_hi\n")
        hpds = self.sga_hpd
        for i in xrange(self.n_cnt):
            f.write(str(self.counts[i])[1:-1]) #strip ellipses
            #hpd = self.sg_ages_model().stats()['Age_%i' % i]['95% HPD interval']
            hpd = hpds[i]
            f.write(", %f,%f,%f, %f, %f,%f\n" % (qages[i][0],qages[i][1], sgages[i][0], sgages[i][1], hpd[0], hpd[1]))
        f.close()
    
    def save(self):
        self.sg2csv()
        pm.utils.coda(self.sga_mcmc())
        pm.utils.coda(self.sa_mcmc())
    
    
    def __str__(self):
        s1 = "\n---------------\nSample %s\n" % self.name
        s2 = "Pooled   age  %.2f\t+- %.2f\n" % self.pooled_age   
        s3 = "MCMC age      %.2f\t+- %.2f\t" % self.sa + str(self.sa_hpd) + \
                        '\tp = %.3f\n' % self.sa_discrepancy()
        return s1+s2+s3
